# -*- coding: utf-8 -*-
"""
拼接图片
"""
import sys
import argparse
import cv2
import os
import random
import copy
import numpy as np
import collections

import rotate
import resize
import flip
import crop
import voc_xml
import utils
from voc_xml import CreateXML
import xml.etree.ElementTree as ET


def confine(value, v_min, v_max):
    """
    值的边界限制
    Args:
        value:输入值
        v_min,v_max:最大最小边界
    return:
        value:限制值
    """
    value = v_min if value < v_min else value
    value = v_max if value > v_max else value
    return value


def random_choose_file(file_dir):
    """
    目录下随机选择一个文件
    Args:
        dir:目录
    return:
        filename:随机选择的文件名
    """
    filename = None
    for root, dirs, files in os.walk(file_dir):
        index = random.randint(0, len(files)-1)
        filename = files[index]
    return filename


def get_xml_tree(xml_file):
    """
    获取xml tree
    Args:
        xml_file: xml文件路径
    return:
        tree:xml tree
    """
    tree = ET.parse(xml_file)
    return tree


def mosaic_img(img, sub_img, start_row, start_col):
    """
    嵌入子图
    Args:
        img:大图
        part_img:待嵌入子图
        start_row,start_col:子图嵌入起始行列
    return:
        img:嵌入结果图
    """
    if sub_img is None:
        return img
    rows, cols, _ = sub_img.shape
    img[start_row:start_row+rows, start_col:start_col+cols] = sub_img
    return img


def translational_box(box, start_row, start_col):
    """
    平移box坐标
    Args:
        box:边框坐标[xmin,ymin,xmax,ymax]
        start_row,start_col:子图嵌入起始行列
    return:
        trans_box:平移后边框坐标
    """
    trans_box = [box[0]+start_col, box[1]+start_row,
                 box[2]+start_col, box[3]+start_row]
    return trans_box


def transform_box(box, augment_options):
    """
    目标框坐标转换
    Args:
        box:目标框[xmin,ymin,xmax,ymax]
        transforms:转换操作[{'opt':'rotate','cterxy':[],'imgwh':[],'rot_angle':0,'randomRotation':False,\
                            'randomAngleRange':[0,360],'scale':1.0,'correction':True,'bk_imgs_dir':'xxx'},
                               {'opt':'crop','crop_type':RANDOM_CROP,'dsize':(0,0),'top_left_x':0,'top_left_y':0,'fw':0.5,'fh':0.7,'random_wh':False ,'iou_thr':0.5},
                               {'opt':'flip','flip_type':-1,'random_flip':True,'imgwh':[]},
                               {'opt':'resize','fx':0.5,'fy':0.5,'dsize':(0,0),'imgwh':[]}]
    return:
        transformed_box:转换后目标框坐标[xmin,ymin,xmax,ymax]        
    """
    transformed_box = box
    for augment_type, options in augment_options.items():
        if [0, 0, 0, 0] == transformed_box:
            break
        if 'rotate' == augment_type:
            #box,cterxy,imgwh,rot_angle,scale=1.0,correction=True
            tmp_box = rotate.rot_box(transformed_box, options['cterxy'],options['imgwh'],
                                     options['rot_angle'], options['scale'], options['correction'])
            imgw, imgh = options['imgwh'][0], options['imgwh'][1]
            transformed_box = [confine(tmp_box[0], 0, imgw-1),
                               confine(tmp_box[1], 0, imgh-1),
                               confine(tmp_box[4], 0, imgw-1),
                               confine(tmp_box[5], 0, imgh-1)]
        elif 'crop' == augment_type:
            transformed_box = crop.crop_box(transformed_box, options['left'],
                                            options['top'], options['crop_w'],
                                            options['crop_h'], options['iou_thr'])
        elif 'flip' == augment_type:
            transformed_box = flip.flip_box(transformed_box, options['imgwh'][0],
                                            options['imgwh'][1], options['flip_type'])
        elif 'resize' == augment_type:
            transformed_box = resize.resize_box(transformed_box, options['scale_w'],
                                                options['scale_h'], options['left'], options['top'])
    return transformed_box


def transform_xml(part_xml_tree, createdxml, augment_options, start_row, start_col):
    """
    将子图的标注框添加到总图的xml中
    Args:
        part_xml_tree:子图xml ET.parse()
        createdxml:总图创建的xml CreateXML对象
        transforms:转换操作
        start_row,start_col:子图嵌入起始行列
    return:
        createdxml: 总图创建的xml CreateXML对象   
    """
    root = part_xml_tree.getroot()
    for obj in root.iter('object'):
        obj_name = obj.find('name').text
        xml_box = obj.find('bndbox')
        xmin = int(xml_box.find('xmin').text)
        ymin = int(xml_box.find('ymin').text)
        xmax = int(xml_box.find('xmax').text)
        ymax = int(xml_box.find('ymax').text)
        box = transform_box([xmin, ymin, xmax, ymax], augment_options)
        if (box[0] >= box[2]) or (box[1] >= box[3]):
            continue
        '''抑制掉面积小于15的小目标 4x4'''
        if (box[2]-box[0]) * (box[3]-box[1]) < 15:
            continue
        box = translational_box(box, start_row, start_col)
        createdxml.add_object_node(obj_name, box[0], box[1], box[2], box[3])
    return createdxml

        
def transform_img(src_img, augment_options):
    """
    图像变换
    Args:
        src_img:源图片
        transforms:转换操作[{'opt':'rotate','cterxy':[],'imgwh':[],'rot_angle':0,'randomRotation':False,\
                            'randomAngleRange':[0,360],'scale':1.0,'correction':True,'bk_imgs_dir':'xxx'},
                               {'opt':'crop','crop_type':RANDOM_CROP,'dsize':(0,0),'top_left_x':0,'top_left_y':0,'fw':0.5,'fh':0.7,'random_wh':False ,'iou_thr':0.5},
                               {'opt':'flip','flip_type':-1,'random_flip':True,'imgwh':[]},
                               {'opt':'resize','fx':0.5,'fy':0.5,'dsize':(0,0),'imgwh':[]}]
    return：
        transformed_img:变换后的图片
        certain_transforms:实际变换操作参数
    """
    transform_options = copy.deepcopy(augment_options)
    actual_trans_options = collections.OrderedDict()
    imgh, imgw, depth = src_img.shape
    img_size = [imgw, imgh]
    transformed_img = src_img
    augment_type = random.choice(list(transform_options))
    options = transform_options[augment_type]
    options['imgwh'] = img_size  # 每一种操作的输入图片宽高

    augment_type = 'flip'  # for dubug
    options = transform_options[augment_type]
    options['imgwh'] = img_size  # 每一种操作的输入图片宽高

    if 'rotate' == augment_type:
        background_img = cv2.imread(os.path.join(options['background_img_dir'],
                                                 random_choose_file(options['background_img_dir'])))
        center_xy = [int(imgw/2), int(imgh/2)]
        rot_angle = options['rot_angle']
        if options['randomRotation']:
            rot_angle = random.randint(options['randomAngleRange'][0], options['randomAngleRange'][1])
        transformed_img = rotate.rot_img_and_padding(transformed_img,
                                                     background_img, center_xy, rot_angle, options['scale'])
        options['cterxy'] = center_xy
        options['rot_angle'] = rot_angle
    elif 'crop' == augment_type:
        scale_w = options['scale_w']
        scale_h = options['scale_h']
        crop_imgw, crop_imgh = options['cropped_size']
        if (0, 0) == options['cropped_size'] and not options['random_wh']:
            crop_imgw = int(options['imgwh'][0] * options['scale_w'])
            crop_imgh = int(options['imgwh'][1] * options['scale_h'])
        elif options['random_wh']:
            crop_imgw = int(options['imgwh'][0] * (scale_w + random.random() * (1 - scale_w)))
            crop_imgh = int(options['imgwh'][1] * (scale_h + random.random() * (1 - scale_h)))

        if 'CENTER_CROP' == options['crop_type']:
            left, top = int(imgw / 2 - crop_imgw / 2), \
                        int(imgh / 2 - crop_imgh / 2)
        elif 'RANDOM_CROP' == options['crop_type']:
            left, top = random.randint(0, imgw - crop_imgw - 1),\
                        random.randint(0, imgh - crop_imgh - 1)
        elif 'FIVE_CROP':
            left, top = options['left'], options['top']

        transformed_img = crop.crop_img(transformed_img, left, top, crop_imgw, crop_imgh)
        if transformed_img is None:
            return None, None
        imgwh = [crop_imgw, crop_imgh]
        options['left'], options['top'] = left, top
        options['crop_w'], options['crop_h'] = crop_imgw, crop_imgh
    elif 'flip' == augment_type:
        flip_type = options['flip_type']
        if options['random_flip']:
            flip_type = random.randint(-1, 1)
        transformed_img = flip.flip_img(transformed_img, flip_type)
        options['flip_type'] = flip_type
    else:
        print("unknown augment type: {}".format(augment_type))
    actual_trans_options[augment_type] = options

    # transformed_size: (height, width)
    transformed_size = (transformed_img.shape[1], transformed_img.shape[0])
    if transformed_size != transform_options['resize']['size']:
        transformed_img, left, top = resize.resize_img(transformed_img,
                                       transform_options['resize']['size'],
                                       transform_options['resize']['scale_w'],
                                       transform_options['resize']['scale_h'],
                                       transform_options['resize']['fixed_ratio'])
        transform_options['resize']['left'] = left
        transform_options['resize']['top'] = top
        actual_trans_options['resize'] = transform_options['resize']
    return transformed_img, actual_trans_options


def mosaic_img_xml(img, sub_img, createdxml, sub_xml_tree, sub_img_size,
                   augment_option, start_row, start_col):
    """
    子图和xml嵌入
    Args:
        img:总图
        part_img:嵌入图
        createdxml：总图创建的xml CreateXML对象
        part_xml_tree：嵌入图xml,ET.parse()
        transforms:转换操作
        start_row,start_col:子图嵌入起始行列
    return:
        img:总图
        createdxml:总图创建的xml CreateXML对象
    """
    transformed_img, transform_options = transform_img(sub_img, augment_option)
    if transformed_img is None or transform_options is None:
        return img, createdxml
    img = mosaic_img(img, transformed_img, start_row, start_col)
    createdxml = transform_xml(sub_xml_tree, createdxml, transform_options, start_row, start_col)
    return img, createdxml


def generate_img_xml(img_save_name, imgw, imgh,
                     part_imgw, part_imgh, augment_option, src_dir_groups):
    """
    生成拼接图和拼接xml
    Args:
        img_save_name:
        imgw,imgh:生成总图宽高
        transforms:转换操作
        #imgs_dir:图源目录
        #xmls_dir:图源对应的xml目录
        imgs_xmls_dirs:源图片、xml、背景图输入目录（支持多组）
    return:
        img:总图
        createdxml：总图创建的xml,ET.parse()
    """
    createdxml = CreateXML(img_save_name, imgw, imgh, 3)
    img = np.zeros((imgh, imgw, 3), dtype=np.uint8) + 255
    
    part_cols = int(imgw / part_imgw)
    part_rows = int(imgh / part_imgh)

    for row in range(part_rows):
        for col in range(part_cols):
            start_row = row * part_imgh
            start_col = col * part_imgw
            
            # 随机选择图片组
            src_dir = src_dir_groups[random.randint(0, len(src_dir_groups)-1)]
            imgs_dir = os.path.join(src_dir, "JPEGImages")
            xmls_dir = os.path.join(src_dir, "Annotations")
            background_img_dir = os.path.join(src_dir, "background_image")
            for augment_type, option in augment_option.items():
                if augment_type == 'rotate':
                    option['background_img_dir'] = background_img_dir
            
            sub_img_name = random_choose_file(imgs_dir)
            if sub_img_name is not None:
                sub_img_full_path = os.path.join(imgs_dir, sub_img_name)
            else:
                print("image {} not found. ".format(sub_img_name))
                continue
            sub_img = cv2.imread(sub_img_full_path)
            
            sub_xml_file_full_path = os.path.join(xmls_dir, sub_img_name.split('.')[0]+'.xml')
            sub_xml_tree = get_xml_tree(sub_xml_file_full_path)
            sub_img_size = (part_imgw, part_imgh)
            img, createdxml = mosaic_img_xml(img, sub_img, createdxml, sub_xml_tree,
                                             sub_img_size, augment_option, start_row, start_col)
    return img, createdxml
            

def generate_img_xml_from_dir(src_dir_groups, imgs_save_dir, xmls_save_dir, name_suffix,
                              count, imgw, imgh, sub_imgw, sub_imgh, augment_option):
    """
    批量拼接图片和xml
    Args:
        #imgs_dir,xmls_dir:源图片和xml路径
        imgs_xmls_dirs:源图片、xml、背景图输入目录（支持多组）
        imgs_save_dir,xmls_save_dir：图片和xml保存路径
        name_suffix: 处理完成的图片、xml的命名标识
        count:生成图片数量
        imgw,imgh:目标拼接图片宽高
        part_imgw,part_imgh:拼接子图宽高
        transforms:转换操作[{'opt':'rotate','cterxy':[],'imgwh':[],'rot_angle':0,'randomRotation':False,\
                            'randomAngleRange':[0,360],'scale':1.0,'correction':True,'bk_imgs_dir':'xxx'},
                        {'opt':'crop','crop_type':RANDOM_CROP,'dsize':(0,0),'top_left_x':0,'top_left_y':0,'fw':0.5,'fh':0.7,'random_wh':False ,'iou_thr':0.5},
                        {'opt':'flip','flip_type':-1,'random_flip':True,'imgwh':[]},
                        {'opt':'resize','fx':0.5,'fy':0.5,'dsize':(0,0),'imgwh':[]}]        
    """

    for n in range(count):
        img_save_name = name_suffix+'_'+str(n)+'.jpg'
        img, createdxml = generate_img_xml(img_save_name, imgw, imgh,
                                           sub_imgw, sub_imgh, augment_option, src_dir_groups)
        cv2.imwrite(os.path.join(imgs_save_dir, img_save_name), img)
        createdxml.save_xml(xmls_save_dir, img_save_name.split('.')[0]+'.xml')
        if n % 5 == 0:
            print("Complete: {}, Total: {}".format(n, count))
        
    
def data_augment(root_dir, save_dir):
    """
    crop, flip, rotate and mosaic
    :param root_dir:
    :param save_dir:
    :param options:
    :return: None
    """

    groups = ["voc2012", "voc2012"]

    src_dir_groups = []
    for group in groups:
        src_dir = os.path.join(root_dir, group)
        src_dir_groups.append(src_dir)

    imgs_save_dir = os.path.join(save_dir, "JPEGImages")
    if not os.path.exists(imgs_save_dir):
        os.makedirs(imgs_save_dir)
    xmls_save_dir = os.path.join(save_dir, "Annotations")
    if not os.path.exists(xmls_save_dir):
        os.makedirs(xmls_save_dir)

    # 命名标识
    name_suffix = 'mosaic_multi4x3'

    # 拼接图片数量
    count = 100

    # 每张拼接图的大小
    img_w, img_h = 800, 600

    sub_img_w, sub_img_h = int(img_w/4), int(img_h/3)

    # crop_type: RANDOM_CROP, CENTER_CROP, FIVE_CROP
    augment_option = {"crop": {'scale_w': 0.5,
                               'scale_h': 0.5,
                               'crop_type': 'RANDOM_CROP',
                               'cropped_size': (sub_img_w, sub_img_h),
                               'left': 0,
                               'top': 0,
                               'crop_w': 0,
                               'crop_h': 0,
                               'random_wh': False,
                               'iou_thr': 0.5,
                               'imgwh': []},
                      "flip":   {'flip_type': -1,
                                 'random_flip': True,
                                 'imgwh': []},
                      "rotate": {'cterxy': [],
                                 'imgwh': [],
                                 'rot_angle': 0,
                                 'randomRotation': True,
                                 'randomAngleRange': [0, 360],
                                 'scale': 1.0,
                                 'correction': True,
                                 'background_img_dir': ''},
                      "resize": {'scale_w': 0.5,
                                 'scale_h': 0.5,
                                 'fixed_ratio': False,
                                 'size': (sub_img_w, sub_img_h),
                                 'left': 0,
                                 'top': 0,
                                 'imgwh': []}}

    generate_img_xml_from_dir(src_dir_groups, imgs_save_dir, xmls_save_dir, name_suffix,
                              count, img_w, img_h, sub_img_w, sub_img_h, augment_option)


if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
        SAVE_DIR = sys.argv[2]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data"),
                            help="VOC dataset root directory")
        PARSER.add_argument("--save_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/masic"),
                            help="directory where augmented data should be saved.")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
        SAVE_DIR = ARGS.save_dir
    data_augment(ROOT_DIR, SAVE_DIR)
